package com.maomao.es.vector;

import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.functionscore.FunctionScoreQueryBuilder;
import org.elasticsearch.index.query.functionscore.ScoreFunctionBuilders;
import org.elasticsearch.index.query.functionscore.ScriptScoreFunctionBuilder;
import org.elasticsearch.script.Script;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;
import java.io.IOException;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;


/**
 * @author mao
 * @date 2025-05-30- 14:59
 * @descript 数据的批量插入
 */
@Slf4j
@Service
public class DataImportService {

    @Autowired
    private VectorDocumentRepository repository;

    @Autowired
    private VectorizationService vectorizationService;

    @Autowired
    private RestHighLevelClient restHighLevelClient;

    @Async
    public void batchImport(List<DocumentRequest> documents) {
        List<VectorDocument> vectorDocs = documents.stream()
                .map(doc -> {
                    VectorDocument vdoc = new VectorDocument();
                    vdoc.setTitle(doc.getTitle());
                    vdoc.setContent(doc.getContent());
                    vdoc.setTitleVector(vectorizationService.textToVector(doc.getTitle()));
                    vdoc.setContentVector(vectorizationService.textToVector(doc.getContent()));
                    return vdoc;
                })
                .collect(Collectors.toList());

        repository.saveAll(vectorDocs);
    }

    public void query(float[] queryVector){
        SearchRequest searchRequest = new SearchRequest("vector_docs");
        //构建脚本查询
        Script script = new Script(
                Script.DEFAULT_SCRIPT_TYPE,
                "painless",
                //pic_vector - 搜索字段
                "cosineSimilarity(params.query_vector, 'titleVector') + 1.0",
                new HashMap<String, Object>() {{
                    //searchData.getPicVector() - 需要搜索的向量数组，List<Double> 类型
                    put("query_vector",queryVector);
                }});
        ScriptScoreFunctionBuilder scriptScoreFunctionBuilder = ScoreFunctionBuilders.scriptFunction(script);

        //构建普通查询
        BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery();

        // 只查存在的数据
//        boolQueryBuilder.filter(QueryBuilders.termQuery(FieldGetterUtil.getPropName(searchData::getState), ConstantEnum.ONE.code));
//
//        // 查询名称
//        if (StringUtils.isNotBlank(searchData.getName())) {
//            boolQueryBuilder.filter(QueryBuilders.termQuery(FieldGetterUtil.getPropName(searchData::getName), searchData.getName()));
//        }

        // 构建查询请求
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        FunctionScoreQueryBuilder functionScoreQueryBuilder = QueryBuilders.functionScoreQuery(boolQueryBuilder, scriptScoreFunctionBuilder);
        searchSourceBuilder.query(functionScoreQueryBuilder);
        searchRequest.source(searchSourceBuilder);

        // 执行搜索
        SearchResponse searchResponse = null;
        try {
            searchResponse = restHighLevelClient.search(searchRequest, RequestOptions.DEFAULT);
            searchResponse.getHits().forEach(hit -> {
                Map<String, Object> sourceAsMap = hit.getSourceAsMap();
                log.info("sourceAsMap:{}", sourceAsMap);
                // System.out.println(sourceAsMap);
            });
        } catch (IOException e) {
            e.printStackTrace();
        }

    }
}
